from collections import deque
from typing import Dict, Any, List
import numpy as np
import random

class ReplayMemory:
    """经验回放记忆缓冲池"""
    
    def __init__(self, capacity: int = 10000):
        """初始化记忆池
        
        Args:
            capacity: 记忆池容量
        """
        self.memory = deque(maxlen=capacity)
    
    def push(self, transition: Dict[str, Any]):
        """存储转换样本
        
        Args:
            transition: 包含(state, action, reward, next_state)的字典
        """
        self.memory.append(transition)
    
    def sample(self, batch_size: int) -> List[Dict[str, Any]]:
        """随机采样批次数据
        
        Args:
            batch_size: 批次大小
            
        Returns:
            采样得到的转换样本列表
        """
        return random.sample(self.memory, min(batch_size, len(self.memory)))
    
    def __len__(self) -> int:
        """返回当前存储的样本数"""
        return len(self.memory)


class PrioritizedReplayMemory(ReplayMemory):
    """带优先级的经验回放"""
    
    def __init__(self, capacity: int = 10000, alpha: float = 0.6):
        """初始化优先级记忆池
        
        Args:
            capacity: 记忆池容量
            alpha: 优先级系数(0-1)
        """
        super().__init__(capacity)
        self.priorities = deque(maxlen=capacity)
        self.alpha = alpha
    
    def push(self, transition: Dict[str, Any], priority: float = 1.0):
        """存储带优先级的转换样本
        
        Args:
            transition: 转换样本
            priority: 初始优先级
        """
        super().push(transition)
        self.priorities.append(priority)
    
    def sample(self, batch_size: int, beta: float = 0.4) -> Dict[str, Any]:
        """优先级采样
        
        Args:
            batch_size: 批次大小
            beta: 重要性采样系数
            
        Returns:
            包含样本和对应权重的字典
        """
        probs = np.array(self.priorities) ** self.alpha
        probs /= probs.sum()
        
        indices = np.random.choice(len(self.memory), batch_size, p=probs)
        samples = [self.memory[i] for i in indices]
        
        weights = (len(self.memory) * probs[indices]) ** (-beta)
        weights /= weights.max()
        
        return {
            'samples': samples,
            'indices': indices,
            'weights': weights
        }
    
    def update_priorities(self, indices: List[int], priorities: List[float]):
        """更新样本优先级
        
        Args:
            indices: 要更新的样本索引
            priorities: 新的优先级值
        """
        for idx, priority in zip(indices, priorities):
            self.priorities[idx] = priority